package core

import (
	"context"
	"encoding/binary"
	"fmt"
	"github.com/kataras/golog"
	"golang.org/x/net/icmp"
	"golang.org/x/net/ipv4"
	"net"
	"ntrace/enviroments"
	"ntrace/utils/routines"
	"time"
)

type ICMPReceiver struct {
	BaseReceiver
}

func NewICMPReceiver(ctx context.Context, task *TraceTask, stat *TraceStat) *ICMPReceiver {
	return &ICMPReceiver{
		BaseReceiver: BaseReceiver{
			Ctx:  ctx,
			Name: "ICMPReceiver",
			Task: task,
			Stat: stat,
		},
	}
}

func (r *ICMPReceiver) StartReceiver() error {
	routines.GoWithRecover("ICMPReceiver", r.doReceive)
	return nil
}

func (r *ICMPReceiver) doReceive() {
	conn, err := OpenReceiverIP4PacketConn("ip4:icmp", "0.0.0.0") // 本地地址
	if err != nil {
		return
	}
	r.Conn = conn

	pid := enviroments.PID
	filter := CreateBPFICMP(uint16(pid), net.ParseIP(r.Task.DstIP))
	err = SetRawConnBPF(conn, filter)
	if err != nil {
		fmt.Println("unable to open network connection")
		return
	}

	payloadBuf := AllocRxBuf(r.Task.MsgLen)
	for {
		select {
		case <-r.Ctx.Done():

			r.StopReceiver()
			return
		default:
			r.receiveOnePkt(pid, payloadBuf)
		}
	}
}

func (r *ICMPReceiver) receiveOnePkt(pid int, payloadBuf []byte) {
	conn := r.Conn
	err := conn.SetReadDeadline(time.Now().Add(100 * time.Millisecond))
	if err != nil {
		return
	}

	h, payload, _, err := conn.ReadFrom(payloadBuf)
	if err != nil {
		return
	}

	devIP := h.Src.To4().String()
	rxTs := time.Now().UnixNano()
	icmpMsg, err := icmp.ParseMessage(1, payload)

	switch icmpMsg.Type {
	case ipv4.ICMPTypeEchoReply:
		echo := icmpMsg.Body.(*icmp.Echo)
		if echo.ID == pid {
			srcIP := h.Dst.To4().String()
			golog.Debugf("pid: %d, srcIP: %s, dstIP: %s, response device: %s, id: %d, seq: %d", pid, srcIP, devIP, devIP, echo.ID, echo.Seq)
			UpdateICMPPktRecord(&r.Stat.PktRecordBucket, devIP, uint16(echo.Seq), devIP, rxTs, true)
		}

	case ipv4.ICMPTypeTimeExceeded:
		originIPPkt := icmpMsg.Body.(*icmp.TimeExceeded).Data

		// is ipv4
		if len(originIPPkt) == 0 || (originIPPkt[0]>>4) != 4 {
			return
		}

		// ip header length
		hLen := (originIPPkt[0] & 0x0F) * 4
		srcIP := net.IPv4(originIPPkt[12], originIPPkt[13], originIPPkt[14], originIPPkt[15])
		dstIP := net.IPv4(originIPPkt[16], originIPPkt[17], originIPPkt[18], originIPPkt[19])

		id := binary.BigEndian.Uint16(originIPPkt[hLen+4 : hLen+6])
		seq := binary.BigEndian.Uint16(originIPPkt[hLen+6 : hLen+8])
		if int(id) == pid {
			golog.Debugf("pid: %d, srcIP: %s, dstIP: %s, response device: %s, id: %d, seq: %d",
				pid, srcIP.To4().String(), dstIP.To4().String(), devIP, id, seq)
			UpdateICMPPktRecord(&r.Stat.PktRecordBucket, dstIP.To4().String(), seq, devIP, rxTs, false)
		}
	}
}

func UpdateICMPPktRecord(pktRecordBucket *PktRecordBucket, target string, seq uint16, device string, rxTs int64, reachDst bool) {
	key := fmt.Sprintf("%s-%d", target, seq)
	pktRecord := &PktRecord{
		Key:      key,
		Device:   device,
		RxTs:     rxTs,
		Received: true,
		ReachDst: reachDst,
	}

	pktRecordBucket.Update(pktRecord)
}
